Broadcasting is a useful NumPy tool that allows us to perform operations between arrays with different shapes, provided that they are compatible with each other in certain ways. To start, we can create an array below and add 5 to it:
In [ ]:
import numpy as np
a = np.array([10, 20, 30, 40])
a + 5
This works even though 5 is not an array; it works like as we would expect, adding 5 to each of the elements in a. This also works if 5 is an array:
In [ ]:
b = np.array([5])
a + b
This takes the single element in b and adds it to each of the elements in a. This won't work for just any b, though; for instance, the following:
b = np.array([5, 6, 7])
a + b
won't work. It does work if a and b are the same shape:
In [ ]:
b = np.array([5, 5, 10, 10])
a + b
What if what we really want is pairwise addition of a, b? Without broadcasting, we could accomplish this by looping:
In [ ]:
b = np.array([1, 2, 3, 4, 5])
In [ ]:
result = np.empty((5, 4), dtype=np.int32)
for row, valb in enumerate(b):
for col, vala in enumerate(a):
result[row, col] = vala + valb
result
We can also do this by manually repeating the arrays to the proper shape for the result, using np.tile. This avoids the need to manually loop:
In [ ]:
aa = np.tile(a, (5, 1))
aa
In [ ]:
# Turn b into a column array, then tile it
bb = np.tile(b.reshape(5, 1), (1, 4))
bb
In [ ]:
aa + bb
We can also do this using broadcasting, which is where NumPy implicitly repeats the array without using additional memory. With broadcasting, NumPy takes care of repeating for you, provided dimensions are "compatible". This works as:
In [ ]:
a.shape
In [ ]:
b.shape
Right now, they have the same number of dimensions, 1, but that dimension is incompatible. We can solve this by appending a dimension using np.newaxis when indexing:
In [ ]:
bb = b[:, np.newaxis]
bb.shape
In [ ]:
a + bb
This can be written more directly in one line:
In [ ]:
a + b[:, np.newaxis]
This also works 2D and 1D, etc.:
In [ ]:
x = np.array([1, 2])
y = np.array([3, 4, 5])
z = np.array([6, 7, 8, 9])
In [ ]:
d_2d = x[:, np.newaxis]**2 + y**2
In [ ]:
d_2d.shape
In [ ]:
d_3d = d_2d[..., np.newaxis] + z**2
In [ ]:
d_3d.shape
Or in one line:
In [ ]:
h = x[:, np.newaxis, np.newaxis]**2 + y[np.newaxis, :, np.newaxis]**2 + z**2
We can see this one-line result has the same shape and same values as the other multi-step calculation.
In [ ]:
h.shape
In [ ]:
np.all(h == d_3d)
Broadcasting is often useful when you want to do calculations with coordinate values, which are often given as 1D arrays corresponding to positions along a particular array dimension. For example, taking range and azimiuth values for radar data (1D separable polar coordinates) and converting to x,y pairs relative to the radar location.
In [ ]:
# Starting data
pressure = np.array([1000, 850, 500, 300])
temps = np.linspace(20, 30, 24).reshape(4, 3, 2)
#
# YOUR CALCULATION HERE
#
In [ ]:
# %load solutions/broadcasting.py
When working with arrays of data, loops over the individual array elements is a fact of life. However, for improved runtime performance, it is important to avoid performing these loops in Python as much as possible, and let NumPy handle the looping for you. Avoiding these loops frequently, but not always, results in shorter and clearer code as well.
In [ ]:
a = np.linspace(0, 20, 6)
a
We can calculate the forward difference for this array with a manual loop as:
In [ ]:
d = np.zeros(a.size - 1)
for i in range(len(a) - 1):
d[i] = a[i + 1] - a[i]
d
It would be nice to express this calculation as a loop, if possible. To see how to go about this, let's condsider the values that are involved in calculating d[i], a[i+1] and a[i]. The values over the loop iterations are:
| i | a[i+1] | a[i] |
| 0 | 4 | 0 |
| 1 | 8 | 4 |
| 2 | 12 | 8 |
| 3 | 16 | 12 |
| 4 | 20 | 16 |
We can express the series of values for a[i+1] then as:
In [ ]:
a[1:]
and a[i] as:
In [ ]:
a[:-1]
This means that we can express the forward difference as:
In [ ]:
a[1:] - a[:-1]
It should be noted that using slices in this way returns only a view on the original array. This means not only can you use the slices to modify the original data (even accidentally), but that this is also a quick operation that does not involve a copy and does not bloat memory usage.
A finite difference estimate of the 2nd derivative is given by: $$f''(x) = 2 f_i - f_{i+1} - f_{i-1}$$
(we're ignoring $\Delta x$ here)
a (using slices)What values should we be expecting to get for the 2nd derivative?
In [ ]:
# %load solutions/vectorized_diff.py
In [ ]:
temps = np.round(20 + np.random.randn(10) * 5, 1)
temps
Let's start by writing a loop to take a 3-point running mean of the data. We'll do this by iterating over all the point in the array and average the 3 points centered on that point. We'll simplify the problem by avoiding dealing with the cases at the edges of the array.
In [ ]:
avg = np.zeros_like(temps)
# We're just ignoring the edge effects here
for i in range(1, len(temps) - 1):
sub = temps[i - 1:i + 2]
avg[i] = sub.mean()
In [ ]:
avg
As with the case of doing finite differences, we can express this using slices of the original array:
In [ ]:
# i - 1 i i + 1
(temps[:-2] + temps[1:-1] + temps[2:]) / 3
Another option to solve this is not using slicing but by using a powerful numpy tool: as_strided. This tool can result in some odd behavior, so take care when using--the tradeoff is that this can be used to do some powerful operations. What we're doing here is altering how NumPy is interpreting the values in the memory that underpins the array. So for this array:
In [ ]:
temps
we can create a view of the array with a new, bigger shape, with rows made up of overlapping values. We do this by specifying a new shape of 8x3, one row for each of the length 3 blocks we can fit in the original 1D array of data. We then use the strides argument to control how numpy walks between items in each dimension. The last item in the strides tuple is just as normal--it says that the number of bytes to walk between items is just the size of an item. (Increasing this would skip items.) The first item says that when we go to a new, in this case row, only advance the size of a single item. This is what gives us overlapping rows.
In [ ]:
block_size = 3
new_shape = (len(temps) - block_size + 1, block_size)
bytes_per_item = temps.dtype.itemsize
temps_strided = np.lib.stride_tricks.as_strided(temps,
shape=new_shape,
strides=(bytes_per_item, bytes_per_item))
temps_strided
Now that we have this view of the array with the rows representing overlapping blocks, we can operate across the rows with mean and the axis=-1 argument to get our running average:
In [ ]:
temps_strided.mean(axis=-1)
It should be noted that there are no copies going on here, so if we change a value at a single indexed location, the change actually shows up in multiple locations:
In [ ]:
temps_strided[0, 2] = 2000
temps_strided
Another operation that crops up when slicing and dicing data is trying to identify a set of indexes, along a particular axis, within a larger multidimensional array. For instance, say we have a 3D array of temperatures, and want to identify the location of the $-10^oC$ isotherm within each column:
In [ ]:
pressure = np.linspace(1000, 100, 25)
temps = np.random.randn(25, 30, 40) * 3 + np.linspace(25, -100, 25).reshape(-1, 1, 1)
NumPy has the function argmin() which returns the index of the minium value. We can use this to find the minimum absolute difference between the value and -10:
In [ ]:
# Using axis=0 to tell it to operate along the pressure dimension
inds = np.argmin(np.abs(temps - -10), axis=0)
inds
In [ ]:
inds.shape
Great! We have an array representing the index of the point closest to $-10^oC$ in each column of data. We could use this to look up into our pressure coordinates to find the pressure level for each column:
In [ ]:
pressure[inds]
How about using that to find the actual temperature value that was closest?
In [ ]:
temps[inds, :, :].shape
Unfortunately, this replaced the pressure dimension (size 25) with the shape of our index array (30 x 40), giving us a 30 x 40 x 30 x 40 array (imagine what would have happened with real data!). One solution here would be to loop:
In [ ]:
output = np.empty(inds.shape, dtype=temps.dtype)
for (i, j), val in np.ndenumerate(inds):
output[i, j] = temps[val, i, j]
output
Of course, what we really want to do is avoid the explicit loop. Let's temporarily simplify the problem to a single dimension. If we have a 1D array, we can pass a 1D array of indices (a full) range, and get back the same as the original data array:
In [ ]:
pressure[np.arange(pressure.size)]
In [ ]:
np.all(pressure[np.arange(pressure.size)] == pressure)
We can use this to select all the indices on the other dimensions of our temperature array. We will also need to use the magic of broadcasting to combine arrays of indices across dimensions:
Now vectorized solution
In [ ]:
y_inds = np.arange(temps.shape[1])[:, np.newaxis]
x_inds = np.arange(temps.shape[2])
temps[inds, y_inds, x_inds]
Now let's say we want to find the relative humidity at the -10C isotherm
In [ ]:
np.all(output == temps[inds, y_inds, x_inds])